package collection

import (
	"encoding/binary"
	"hash/fnv"
	"iter"
	"sync"
	"time"

	"gitee.com/gousing/helper/hasher"
	"gitee.com/gousing/helper/jsoner"
)

// MapAs 明确KV类型的Map分段锁数据结构(并发安全)
//   - 使用分段锁控制并发场景，默认32段
//   - 适用于读多写少占比场景，请调整分段数, 建议32或以上
//   - 适用于读少写多均衡场景，请提高分段数，建议32或以下
//   - 适用于读写均衡场景，请提高分段数，建议32
type MapAsSync[K comparable, V any] struct {
	shards   []*mapAsShard[K, V]
	numShard uint32
}
type mapAsShard[K comparable, V any] struct {
	data map[K]V
	mu   sync.RWMutex
}

func NewMapAsSync[K comparable, V any](numShards ...uint32) *MapAsSync[K, V] {
	var numShard uint32 = 32
	if len(numShards) > 0 && numShards[0] > 0 {
		numShard = numShards[0]
	}
	// 分片数大于1时
	if numShard > 1 {
		//key 使用泛型类型，避免编译器优化(类型推断)
		var keyTest K
		var keyAny any = keyTest
		switch keyAny.(type) {
		case bool:
			//如果缓存Key的类型是bool，则分片数固定为1
			numShard = 1
		}
	}
	m := MapAsSync[K, V]{shards: make([]*mapAsShard[K, V], numShard), numShard: numShard}
	var i uint32
	for i = 0; i < numShard; i++ {
		m.shards[i] = &mapAsShard[K, V]{
			data: make(map[K]V, 0),
			mu:   sync.RWMutex{},
		}
	}
	return &m
}

const (
	hashFnvOffset32 uint32 = 2166136261
	hashFnvPrime32  uint32 = 16777619
)

// shardSync 计算散列值时，key 使用泛型类型，避免编译器优化(类型推断)
func (m *MapAsSync[K, V]) shardSync(key K) *mapAsShard[K, V] {
	if m.numShard == 1 {
		return m.shards[0]
	}
	// shardSync 计算散列值时，key 使用泛型类型，避免编译器优化(类型推断)
	var keyAny any = key

	if keyAny == nil {
		return m.shards[hashFnvOffset32%m.numShard]
	}
	switch keyVal := keyAny.(type) {
	case string:
		hash := hashFnvOffset32
		for i := 0; i < len(keyVal); i++ {
			hash ^= uint32(keyVal[i])
			hash *= hashFnvPrime32
		}
		return m.shards[hash%m.numShard]
	case int:
		h := fnv.New32a()
		if is64Bit {
			k64 := int64(keyVal)
			binary.Write(h, binary.BigEndian, &k64)
		} else {
			k32 := int32(keyVal)
			binary.Write(h, binary.BigEndian, &k32)
		}
		return m.shards[h.Sum32()%m.numShard]
	case int8:
		h := fnv.New32a()
		binary.Write(h, binary.BigEndian, &keyVal)
		return m.shards[h.Sum32()%m.numShard]
	case int16:
		h := fnv.New32a()
		binary.Write(h, binary.BigEndian, &keyVal)
		return m.shards[h.Sum32()%m.numShard]
	case int32:
		h := fnv.New32a()
		binary.Write(h, binary.BigEndian, &keyVal)
		return m.shards[h.Sum32()%m.numShard]
	case int64:
		h := fnv.New32a()
		binary.Write(h, binary.BigEndian, &keyVal)
		return m.shards[h.Sum32()%m.numShard]
	case uint:
		h := fnv.New32a()
		if is64Bit {
			u64 := uint64(keyVal)
			_, _ = h.Write([]byte{byte(u64 >> 56), byte(u64 >> 48), byte(u64 >> 40), byte(u64 >> 32), byte(u64 >> 24), byte(u64 >> 16), byte(u64 >> 8), byte(u64)})
		} else {
			u32 := uint32(keyVal)
			_, _ = h.Write([]byte{byte(u32 >> 24), byte(u32 >> 16), byte(u32 >> 8), byte(u32)})
		}
		return m.shards[h.Sum32()%m.numShard]
	case uint8:
		h := fnv.New32a()
		_, _ = h.Write([]byte{byte(keyVal)})
		return m.shards[h.Sum32()%m.numShard]
	case uint16:
		h := fnv.New32a()
		_, _ = h.Write([]byte{byte(keyVal >> 8), byte(keyVal)})
		return m.shards[h.Sum32()%m.numShard]
	case uint32:
		h := fnv.New32a()
		_, _ = h.Write([]byte{byte(keyVal >> 24), byte(keyVal >> 16), byte(keyVal >> 8), byte(keyVal)})
		return m.shards[h.Sum32()%m.numShard]
	case uint64:
		h := fnv.New32a()
		_, _ = h.Write([]byte{byte(keyVal >> 56), byte(keyVal >> 48), byte(keyVal >> 40), byte(keyVal >> 32), byte(keyVal >> 24), byte(keyVal >> 16), byte(keyVal >> 8), byte(keyVal)})
		return m.shards[h.Sum32()%m.numShard]
	case float32:
		h := fnv.New32a()
		binary.Write(h, binary.BigEndian, &keyVal)
		return m.shards[h.Sum32()%m.numShard]
	case float64:
		h := fnv.New32a()
		binary.Write(h, binary.BigEndian, &keyVal)
		return m.shards[h.Sum32()%m.numShard]
	case time.Time:
		if keyVal.IsZero() {
			return m.shards[hashFnvOffset32%m.numShard]
		}
		h := fnv.New32a()
		k64 := keyVal.UnixNano()
		binary.Write(h, binary.BigEndian, &k64)
		return m.shards[h.Sum32()%m.numShard]
	case time.Duration:
		if keyVal == 0 {
			return m.shards[hashFnvOffset32%m.numShard]
		}
		h := fnv.New32a()
		k64 := keyVal.Nanoseconds()
		binary.Write(h, binary.BigEndian, &k64)
		return m.shards[h.Sum32()%m.numShard]
	case bool:
		// Key 为 bool 类型时，分片数固定为1，不会进行此运算
		if keyVal {
			return m.shards[1%m.numShard]
		} else {
			return m.shards[0%m.numShard]
		}
	case complex64:
		if keyVal == 0 {
			return m.shards[hashFnvOffset32%m.numShard]
		}
		h := fnv.New32a()
		binary.Write(h, binary.LittleEndian, real(keyVal))
		binary.Write(h, binary.LittleEndian, imag(keyVal))
		return m.shards[h.Sum32()%m.numShard]
	case complex128:
		if keyVal == 0 {
			return m.shards[hashFnvOffset32%m.numShard]
		}
		h := fnv.New32a()
		binary.Write(h, binary.LittleEndian, real(keyVal))
		binary.Write(h, binary.LittleEndian, imag(keyVal))
		return m.shards[h.Sum32()%m.numShard]
	default:
		return m.shards[hasher.Fnv32Any(key)%m.numShard]
	}
}

// Set 添加键值，单层替换方式
func (m *MapAsSync[K, V]) Set(key K, val V) *MapAsSync[K, V] {
	m.shardSync(key).set(key, val)
	return m
}

// Merge 合并map到*MapShard, 单层替换方式合并
func (m *MapAsSync[K, V]) Merge(maps map[K]V) *MapAsSync[K, V] {
	if len(maps) == 0 {
		return m
	}
	for key, val := range maps {
		m.shardSync(key).set(key, val)
	}
	return m
}

// Has 判断是否存在key
func (m *MapAsSync[K, V]) Has(key K) bool {
	return m.shardSync(key).has(key)
}

// GetE 取得map[K]V数据集的元素值
func (m *MapAsSync[K, V]) GetE(key K) (val V, ok bool) {
	return m.shardSync(key).get(key)
}

// GetD 取得map[K]V数据集的元素值，不存在时返回默认值
func (m *MapAsSync[K, V]) GetD(key K, defaultVal V) (val V) {
	if v, ok := m.shardSync(key).get(key); ok {
		return v
	}
	return defaultVal
}

// Delete 删除map[K]V数据集的元素值
func (m *MapAsSync[K, V]) Delete(key K) *MapAsSync[K, V] {
	m.shardSync(key).delete(key)
	return m
}

// Size 获取数据集大小
func (m *MapAsSync[K, V]) Size() (size int) {
	for _, sd := range m.shards {
		size += sd.size()
	}
	return
}

// Clear 清空所有数据
func (m *MapAsSync[K, V]) Clear() {
	for _, sd := range m.shards {
		sd.clear()
	}
}

// Collects 取得map[k]v数据集的浅克隆副本  map[string]any
func (m *MapAsSync[K, V]) Collects() (data map[K]V) {
	data = make(map[K]V)
	for _, sd := range m.shards {
		for key, val := range sd.mapData() {
			data[key] = val
		}
	}
	return data
}

// Keys 取得map键 数据集的键名副本 []string
func (m *MapAsSync[K, V]) Keys() (data []K) {
	data = make([]K, 0)
	for _, sd := range m.shards {
		if sd.size() == 0 {
			continue
		}
		data = append(data, sd.keys()...)
	}
	return data
}

// Values 取得map值 数据集的浅克隆副本 []V
func (m *MapAsSync[K, V]) Values() (data []V) {
	data = make([]V, 0)
	for _, sd := range m.shards {
		if sd.size() == 0 {
			continue
		}
		data = append(data, sd.values()...)
	}
	return data
}

// Iter mak[K]V 元素迭代器
//   - for k,v:=range s.Iter(){}
//   - 注意：迭代顺序不固定, 依据Shard分片+Map[K]V顺序进行迭代
func (m *MapAsSync[K, V]) Iter() iter.Seq2[K, V] {
	return func(yield func(key K, val V) bool) {
		for _, sd := range m.shards {
			if sd.size() == 0 {
				continue
			}
			for k, v := range sd.mapData() {
				if !yield(k, v) {
					return
				}
			}
		}
	}
}

// IterKeys map[K]V key迭代器
//   - for key:=range s.IterKeys(){}
func (m *MapAsSync[K, V]) IterKeys() iter.Seq[K] {
	return func(yield func(key K) bool) {
		for _, sd := range m.shards {
			if sd.size() == 0 {
				continue
			}
			for k := range sd.mapData() {
				if !yield(k) {
					return
				}
			}
		}
	}
}

// IterValues map[K]V value 迭代器
//   - for val:=range s.IterValues(){}
func (m *MapAsSync[K, V]) IterValues() iter.Seq[V] {
	return func(yield func(val V) bool) {
		for _, sd := range m.shards {
			if sd.size() == 0 {
				continue
			}
			for _, v := range sd.mapData() {
				if !yield(v) {
					return
				}
			}
		}
	}
}

// MarshalJSON 序列化集合接口
func (m *MapAsSync[K, V]) MarshalJSON() ([]byte, error) {
	items := make(map[K]V)
	for _, sd := range m.shards {
		for key, val := range sd.mapData() {
			items[key] = val
		}
	}
	return jsoner.Marshal(items)
}

// UnmarshalJSON 反序列化集合接口
func (m *MapAsSync[K, V]) UnmarshalJSON(data []byte) error {
	items := make(map[K]V)
	err := jsoner.Unmarshal(data, &items)
	if err != nil {
		return err
	}
	m.Clear()
	for key, val := range items {
		m.shardSync(key).set(key, val)
	}
	return nil
}

// mapKvSyncShard 对象驱动
// set 设置Shard map k/v 键值
func (m *mapAsShard[K, V]) set(key K, val V) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.data[key] = val
}

// delete 删除Shard map值
func (m *mapAsShard[K, V]) delete(key K) {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.data, key)
}

// clear 清空Shard map数据
func (m *mapAsShard[K, V]) clear() {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.data = make(map[K]V)
}

// has 判断Shard map是否存在key
func (m *mapAsShard[K, V]) has(key K) bool {
	m.mu.RLock()
	defer m.mu.RUnlock()
	if _, ok := m.data[key]; ok {
		return true
	}
	return false
}

// get 取得Shard map值
func (m *mapAsShard[K, V]) get(key K) (val V, ok bool) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	val, ok = m.data[key]
	return
}

// size 取得Shard map长度
func (m *mapAsShard[K, V]) size() int {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return len(m.data)
}

// mapData 取得Shard map原始数据
func (m *mapAsShard[K, V]) mapData() map[K]V {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.data
}

// keys 取得Shard map键数据集的键名副本 []K
func (m *mapAsShard[K, V]) keys() (keys []K) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	keys = make([]K, len(m.data))
	i := 0
	for key := range m.data {
		keys[i] = key
		i++
	}
	return keys
}

// values 取得Shard map值数据集的浅克隆副本 []V
func (m *mapAsShard[K, V]) values() (values []V) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	values = make([]V, len(m.data))
	i := 0
	for key := range m.data {
		values[i] = m.data[key]
		i++
	}
	return values
}
